#include "pch.h"
#include "DesktopCapturerWin.h"
#include "DesktopCaptureOptions.h"
#include "Utils.h"
#include "Desktop.h"
#include "DesktopFrameWin.h"

namespace {
	const UINT DWM_EC_DISABLECOMPOSITION = 0;
	const UINT DWM_EC_ENABLECOMPOSITION = 1;

	const wchar_t kDwmapiLibraryName[] = L"dwmapi.dll";

} 

DesktopCapturerWin::DesktopCapturerWin(
	const DesktopCaptureOptions& options) {
	if (options.disable_effects()) {
		if (!dwmapi_library_)
			dwmapi_library_ = LoadLibraryW(kDwmapiLibraryName);

		if (dwmapi_library_) {
			composition_func_ = reinterpret_cast<DwmEnableCompositionFunc>(
				GetProcAddress(dwmapi_library_, "DwmEnableComposition"));
		}
	}
}

DesktopCapturerWin::~DesktopCapturerWin() {
	if (desktop_dc_)
		ReleaseDC(NULL, desktop_dc_);
	if (memory_dc_)
		DeleteDC(memory_dc_);

	// Restore Aero.
	if (composition_func_)
		(*composition_func_)(DWM_EC_ENABLECOMPOSITION);

	if (dwmapi_library_)
		FreeLibrary(dwmapi_library_);
}

void DesktopCapturerWin::SetSharedMemoryFactory(
	std::unique_ptr<SharedMemoryFactory> shared_memory_factory) {
	shared_memory_factory_ = std::move(shared_memory_factory);
}

void DesktopCapturerWin::CaptureFrame() {
	int64_t capture_start_time_nanos = Utils::TimeNanos();

	queue_.MoveToNextFrame();

	PrepareCaptureResources();

	if (!CaptureImage()) {
		printf("Failed to capture screen by GDI.");
		callback_->OnCaptureResult(Result::ERROR_TEMPORARY, nullptr);
		return;
	}

	std::unique_ptr<DesktopFrame> frame = queue_.current_frame()->Share();
	frame->set_dpi(DesktopVector(GetDeviceCaps(desktop_dc_, LOGPIXELSX),
		GetDeviceCaps(desktop_dc_, LOGPIXELSY)));
	frame->mutable_updated_region()->SetRect(
		DesktopRect::MakeSize(frame->size()));
	frame->set_capture_time_ms((Utils::TimeNanos() - capture_start_time_nanos) /
		Utils::kNumNanosecsPerMillisec);
	frame->set_capturer_id(DesktopCapturerId::kScreenCapturerWinGdi);

	callback_->OnCaptureResult(Result::SUCCESS, std::move(frame));
}

bool DesktopCapturerWin::GetSourceList(DesktopCapturerWin::SourceList* sources) {
	return Utils::GetScreenList(sources);
}

bool DesktopCapturerWin::SelectSource(SourceId id) {
	bool valid = Utils::IsScreenValid(id, &current_device_key_);
	if (valid)
		current_screen_id_ = id;
	return valid;
}

void DesktopCapturerWin::Start(Callback* callback) {
	callback_ = callback;

	if (composition_func_)
		(*composition_func_)(DWM_EC_DISABLECOMPOSITION);
}

void DesktopCapturerWin::PrepareCaptureResources() {
	std::unique_ptr<Desktop> input_desktop(Desktop::GetInputDesktop());
	if (input_desktop && !desktop_.IsSame(*input_desktop)) {
		if (desktop_dc_) {
			ReleaseDC(NULL, desktop_dc_);
			desktop_dc_ = nullptr;
		}

		if (memory_dc_) {
			DeleteDC(memory_dc_);
			memory_dc_ = nullptr;
		}

		desktop_.SetThreadDesktop(input_desktop.release());

		if (composition_func_) {
			(*composition_func_)(DWM_EC_DISABLECOMPOSITION);
		}
	}

	if (display_configuration_monitor_.IsChanged()) {
		if (desktop_dc_) {
			ReleaseDC(NULL, desktop_dc_);
			desktop_dc_ = nullptr;
		}
		if (memory_dc_) {
			DeleteDC(memory_dc_);
			memory_dc_ = nullptr;
		}
	}

	if (!desktop_dc_) {
		desktop_dc_ = GetDC(nullptr);
		memory_dc_ = CreateCompatibleDC(desktop_dc_);

		queue_.Reset();
	}
}

bool DesktopCapturerWin::CaptureImage() {
	DesktopRect screen_rect =
		Utils::GetScreenRect(current_screen_id_, current_device_key_);
	if (screen_rect.is_empty()) {
		return false;
	}

	DesktopSize size = screen_rect.size();
	if (!queue_.current_frame() ||
		!queue_.current_frame()->size().equals(screen_rect.size())) {

		std::unique_ptr<DesktopFrame> buffer = DesktopFrameWin::Create(
			size, shared_memory_factory_.get(), desktop_dc_);
		if (!buffer) {
			return false;
		}
		queue_.ReplaceCurrentFrame(SharedDesktopFrame::Wrap(std::move(buffer)));
	}
	queue_.current_frame()->set_top_left(
		screen_rect.top_left().subtract(Utils::GetFullscreenRect().top_left()));

	DesktopFrameWin* current = static_cast<DesktopFrameWin*>(
		queue_.current_frame()->GetUnderlyingFrame());
	HGDIOBJ previous_object = SelectObject(memory_dc_, current->bitmap());
	if (!previous_object || previous_object == HGDI_ERROR) {
		return false;
	}

	bool result = (BitBlt(memory_dc_, 0, 0, screen_rect.width(),
		screen_rect.height(), desktop_dc_, screen_rect.left(),
		screen_rect.top(), SRCCOPY | CAPTUREBLT) != FALSE);
	if (!result) {
		printf("BitBlt failed");
	}

	SelectObject(memory_dc_, previous_object);

	return result;
}
